// server/main.go
package main

import (
	"crypto/tls"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"sync"
	"time"

	_ "embed"

	"github.com/gorilla/websocket"
)

//go:embed cert.pem
var certPEM []byte

//go:embed key.pem
var keyPEM []byte

var (
	upgrader = websocket.Upgrader{
		CheckOrigin: func(r *http.Request) bool { return true },
	}
	clients = sync.Map{}
)

func registerHandler(w http.ResponseWriter, r *http.Request) {
	port := r.URL.Query().Get("port")
	if port == "" {
		http.Error(w, "missing port", http.StatusBadRequest)
		return
	}
	subdomain := fmt.Sprintf("%x", time.Now().UnixNano())[:6]
	clients.Store(subdomain, nil)

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]string{
		"subdomain": subdomain,
		"ws_url":    "wss://" + subdomain + ".localhost:8443/tunnel",
	})
}

func tunnelHandler(w http.ResponseWriter, r *http.Request) {
	subdomain := extractSubdomain(r.Host)
	if subdomain == "" {
		http.Error(w, "invalid host", http.StatusBadRequest)
		return
	}
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("upgrade error: %v", err)
		return
	}
	clients.Store(subdomain, conn)
	select {}
}

func proxyHandler(w http.ResponseWriter, r *http.Request) {
	subdomain := extractSubdomain(r.Host)
	if conn, ok := clients.Load(subdomain); ok && conn != nil {
		c := conn.(*websocket.Conn)
		reqData, _ := json.Marshal(map[string]interface{}{
			"method": r.Method,
			"url":    r.URL.String(),
			"header": r.Header,
			"body":   readAll(r.Body),
		})
		c.WriteMessage(websocket.BinaryMessage, reqData)

		_, respData, err := c.ReadMessage()
		if err != nil {
			http.Error(w, "tunnel error", http.StatusBadGateway)
			return
		}

		var resp map[string]interface{}
		json.Unmarshal(respData, &resp)
		for k, v := range resp["header"].(map[string]interface{}) {
			w.Header().Set(k, v.([]interface{})[0].(string))
		}
		w.WriteHeader(int(resp["status"].(float64)))
		w.Write([]byte(resp["body"].(string)))
	} else {
		http.Error(w, "tunnel not found", http.StatusNotFound)
	}
}

func extractSubdomain(host string) string {
	if host == "localhost:8443" {
		return ""
	}
	if i := len(host) - len(".localhost:8443"); i > 0 {
		return host[:i]
	}
	return ""
}

func readAll(r io.Reader) []byte {
	data, _ := io.ReadAll(r)
	return data
}

func main() {
	http.HandleFunc("/register", registerHandler)
	http.HandleFunc("/tunnel", tunnelHandler)
	http.HandleFunc("/", proxyHandler)

	cert, err := tls.X509KeyPair(certPEM, keyPEM)
	if err != nil {
		log.Fatal("Failed to load cert:", err)
	}

	server := &http.Server{
		Addr: ":8443",
		TLSConfig: &tls.Config{
			Certificates: []tls.Certificate{cert},
		},
	}

	log.Println("Tunnel server running on https://localhost:8443")
	log.Fatal(server.ListenAndServeTLS("", ""))
}